#!/usr/bin/env perl

# Copyright 2014 The Souper Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

use warnings;
use strict;
use File::Temp;

sub compiling() {
    foreach my $arg (@ARGV) {
        return 1
            if ($arg =~ /\.c$|\.cpp$|\.CC$|\.c\+\+$|\.cc$|\.cxx$|\.C$|\.c\+$/);
    }
    return 0;
}

sub linkp() {
    foreach my $arg (@ARGV) {
        return 0 if ($arg eq "-S" || $arg eq "-c" || $arg eq "-shared");
    }
    return 1;
}

if ($0 =~ /clang$/) {
    unshift @ARGV, "@LLVM_BINDIR@/clang";
} elsif ($0 =~ /clang\+\+$/) {
    unshift @ARGV, "@LLVM_BINDIR@/clang++";
} else {
    die "Didn't expect sclang to be invoked as '$0'";
}

foreach my $arg (@ARGV) {
    if ($arg eq "-help" || $arg eq "--help") {
	print <<EOF;

sclang / sclang++ emulate clang / clang++ except that they also invoke
Souper.

To avoid the need to mess with build systems, communication with these
programs is via environment variables. Unless noted, the value of the
environment variable is irrelevant: sclang only checks if it is set.

SOUPER_DEBUG -- Print debugging info.

SOUPER_DYNAMIC_PROFILE -- Instrument the compiled program such that,
when run, it will communicate with a running Redis instance in order
to report how many times each optimized code site executes.

SOUPER_DYNAMIC_PROFILE_ALL -- Like SOUPER_DYNAMIC_PROFILE but
instruments each site that is potentially optimizable by Souper, as
opposed to sites that actually get optimized. Should be used in
combination with SOUPER_NO_INFER.

SOUPER_EXPLOIT_BLOCKPCS -- Optimize using blockpcs.

SOUPER_FIRST_OPT / SOUPER_LAST_OPT -- Assign a number to each
optimization found in the current compilation unit and only perform
those that fall within the range [SOUPER_FIRST_OPT
.. SOUPER_LAST_OPT], inclusive. Used for debugging Souper

SOUPER_INFER_INST -- Try to synthesis instructions with InstSynthesis
synthesizer.

SOUPER_ENUMERATIVE_SYNTHESIS_MAX_INSTS -- Number of instructions for
EnumerativeSynthesis synthesizer.

SOUPER_DATAFLOW_PRUNING -- Enable dataflow pruning. This only makes sense
when EnumerativeSynthesis is enabled.

SOUPER_REDIS_PORT -- Mention the port to use the redis-server on.

SOUPER_NO_EXTERNAL_CACHE -- Don't ask the running Redis instance for
cached inferences.

SOUPER_NO_HARVEST_DATAFLOW_FACTS -- Don't query LLVM's bit-level dataflow
analyses when harvesting.

SOUPER_HARVEST_USES -- This option enables the harvest of LLVM uses.

SOUPER_NO_INFER -- Don't run inference at all; used for populating the
cache.

SOUPER_NO_SOUPER -- Don't run Souper at all.

SOUPER_SAVE_BITCODE -- The value of this variable is the name of a
directory where a copy of the bitcode for the compiled file will be
stored.

SOUPER_SKIP_FILES -- Do not apply Souper (as if SOUPER_NO_SOUPER were
set) to any of the comma-separated list of files found in the value of
this variable.

SOUPER_STATIC_PROFILE -- During compilation, communicate with a running
Redis instance in order to record the number of times each optimization
was performed.

SOUPER_USE_ALIVE -- Send queries to Alive2

SOUPER_DOUBLE_CHECK -- Only verify results with Alive2

SOUPER_DISABLE_LLVM_PEEPHOLES -- Disable the invocation of llvm peephole
transformations in compilation.

SOUPER_STATS -- Ask LLVM to print statistics.

SOUPER_TIME_REPORT -- Ask LLVM to print a time report for passes

SOUPER_VERIFY -- Run the module verifier before and after Souper runs

EOF
        exit(-1);
    }
}

my %whitelist = ();

# N.B. have to get environment variables using this function which has
# the side effect of adding the env var to our whitelist
sub getenv($) {
    (my $e) = @_;
    die "oops '$e'" unless $e =~ /^SOUPER_/;
    $whitelist{$e} = 1;
    return undef unless exists($ENV{$e});
    return $ENV{$e};
}

my $bitcodedir = getenv("SOUPER_SAVE_BITCODE");
if (defined $bitcodedir && -d $bitcodedir && fork() == 0) {
    for (my $i=0; $i<scalar(@ARGV); $i++) {
	splice @ARGV, $i, 2 if ($ARGV[$i] eq "-o");
    }
    my $tmp = File::Temp->new(UNLINK => 1, DIR => $bitcodedir)->filename;
    push @ARGV, "-c", "-emit-llvm";
    push @ARGV, "-o", "${tmp}.bc";
    my $ofn = "${tmp}.cmd";
    open OUTF, ">$ofn" or die;
    foreach my $a (@ARGV) {
	print OUTF "$a ";
    }
    print OUTF "\n";
    close OUTF;
    open STDOUT, '>/dev/null';
    open STDERR, '>/dev/null';
    exec @ARGV;
    die "bailing, exec failed";
}

my $souper = 1;
# this environment variable is a comma-separated list of source files that
# souper should avoid processing, for example because they trigger known bugs
my $skip_files = getenv("SOUPER_SKIP_FILES");
if (defined $skip_files) {
    my %skips;
    foreach my $f (split(',', $skip_files)) {
        $skips{$f} = 1;
    }
    foreach my $a (@ARGV) {
        $souper = 0 if ($skips{$a});
    }
}

$souper = 0 if getenv("SOUPER_NO_SOUPER");
$souper = 0 unless compiling();

if ($souper) {
    push @ARGV, (
        "-Xclang", "-load",
        "-mllvm", "-solver-timeout=15",
        );
}

if (getenv("SOUPER_DYNAMIC_PROFILE_ALL")) {
    if ($souper) {
        push @ARGV, ("-Xclang", "@SOUPER_PASS_PROFILE_ALL@");
    }
} else {
    if ($souper) {
        push @ARGV, ("-Xclang", "@SOUPER_PASS@");
    }
}

if (getenv("SOUPER_DEBUG") && $souper) {
    push @ARGV, ("-mllvm", "-souper-debug-level=".getenv("SOUPER_DEBUG"));
}

if (getenv("SOUPER_EXPLOIT_BLOCKPCS") && $souper) {
    push @ARGV, ("-mllvm", "-souper-exploit-blockpcs");
}

if (!getenv("SOUPER_NO_EXTERNAL_CACHE") && $souper) {
    push @ARGV, ("-mllvm", "-souper-external-cache");
}

if (getenv("SOUPER_NO_INFER") && $souper) {
    push @ARGV, ("-mllvm", "-souper-no-infer");
}

if (getenv("SOUPER_FIRST_OPT") && $souper) {
    push @ARGV, ("-mllvm", "-souper-first-opt=".$ENV{"SOUPER_FIRST_OPT"});
}

if (getenv("SOUPER_LAST_OPT") && $souper) {
    push @ARGV, ("-mllvm", "-souper-last-opt=".$ENV{"SOUPER_LAST_OPT"});
}

if (getenv("SOUPER_STATIC_PROFILE") && $souper) {
    push @ARGV, ("-mllvm", "-souper-static-profile");
}

if (getenv("SOUPER_DYNAMIC_PROFILE") ||
    getenv("SOUPER_DYNAMIC_PROFILE_ALL") && $souper) {
    push @ARGV, ("-g", "-mllvm", "-souper-dynamic-profile");
}

if (getenv("SOUPER_USE_ALIVE") && $souper) {
    push @ARGV, ("-mllvm", "-souper-use-alive");
}

if (getenv("SOUPER_DOUBLE_CHECK") && $souper) {
    push @ARGV, ("-mllvm", "-souper-double-check");
}

if (getenv("SOUPER_INFER_INST") && $souper) {
    push @ARGV, ("-mllvm", "-souper-infer-inst");
}

if (getenv("SOUPER_ENUMERATIVE_SYNTHESIS_MAX_INSTS") && $souper) {
    push @ARGV, ("-mllvm", "-souper-enumerative-synthesis-max-instructions=".$ENV{"SOUPER_ENUMERATIVE_SYNTHESIS_MAX_INSTS"});
}

if (getenv("SOUPER_DATAFLOW_PRUNING") && $souper) {
    push @ARGV, ("-mllvm", "-souper-dataflow-pruning");
}

if (getenv("SOUPER_REDIS_PORT") && $souper) {
    push @ARGV, ("-mllvm", "-souper-redis-port=".$ENV{"SOUPER_REDIS_PORT"});
}

if (getenv("SOUPER_STATS") && $souper) {
    push @ARGV, ("-mllvm", "-stats");
}

if (getenv("SOUPER_TIME_REPORT") && $souper) {
    push @ARGV, ("-ftime-report");
}

if (getenv("SOUPER_VERIFY") && $souper) {
    push @ARGV, ("-mllvm", "-souper-verify");
}

if (getenv("SOUPER_NO_HARVEST_DATAFLOW_FACTS") && $souper) {
    push @ARGV, ("-mllvm", "-souper-harvest-dataflow-facts=false");
}

if (getenv("SOUPER_HARVEST_USES") && $souper) {
    push @ARGV, ("-mllvm", "-souper-harvest-uses");
}

if (getenv("SOUPER_DISABLE_LLVM_PEEPHOLES") && $souper) {
    push @ARGV, ("-mllvm", "-disable-all-peepholes");
}

if ((getenv("SOUPER_DYNAMIC_PROFILE") ||
     getenv("SOUPER_DYNAMIC_PROFILE_ALL")) && linkp() && $souper) {
    push @ARGV, ("@PROFILE_LIBRARY@", "@HIREDIS_LIBRARY@");
}

if (getenv("SOUPER_DEBUG") && $souper) {
    foreach my $arg (@ARGV) {
        print STDERR "$arg ";
    }
    print STDERR "\n";
}

foreach my $e (keys %ENV) {
    next unless $e =~ /^SOUPER_/;
    die "unexpected Souper-related environment variable '${e}'" unless $whitelist{$e};
}
    
exec @ARGV;
